# -*- coding: utf-8 -*-
"""
 * @file main.py
 * @author SDK_TEAM
 * @brief
 * @version 0.1
 * @date 2022-11-3
 * Copyright:
 * © 2018 北京灵汐科技有限公司 版权所有。
 * 注意：以下内容均为北京灵汐科技有限公司原创，未经本公司允许，不得转载，否则将视为侵权；对于不遵守此声明或者其他违法使用以下内容者，本公司依法保留追究权。
 * © 2018 Lynxi Technologies Co., Ltd. All rights reserved.
 * NOTICE: All information contained here is, and remains the property of Lynxi. This file can not
 * be copied or distributed without the permission of Lynxi Technologies Co., Ltd.
"""

import argparse
from collections import namedtuple
import pylynchipsdk as sdk
import numpy as np
import threading
import _thread
import struct
from ctypes import *
from typing import cast
import getopt

import os
import sys
import time

sys.path.append("../../../")
from common.python.plugin_utils import *
import common.python.common as common
import common.python.bufferpool as bufferpool
from common.python.common import *
from common.python import image_decoder
from common.python import image_encoder
import ipeParamYolov8

pythonapi.PyCapsule_GetPointer.restype = c_void_p
pythonapi.PyCapsule_GetPointer.argtypes = [py_object, c_char_p]


def main():
    parser = argparse.ArgumentParser(description="yolov8 segment")

    parser.add_argument(
        "input_image",
        metavar="input_image",
        type=str,
        default="../../../data/yolov8_segment_with_image_sample/test.jpeg",
        help="input jpeg file path",
    )
    parser.add_argument(
        "output_image",
        metavar="output_image",
        type=str,
        default="./output.jpeg",
        help="output jpeg file path",
    )
    parser.add_argument(
        "device_id", metavar="device_id", type=int, default=0, help="lynxi device id"
    )
    parser.add_argument(
        "--model",
        metavar="model",
        type=str,
        default="../../../model/yolov8s-seg_b1_h640_w640_c3_iuint8_ofloat32_0000/Net_0",
        help="yolov8 model path",
    )
    parser.add_argument(
        "--post_plugin",
        metavar="post_plugin",
        type=str,
        default="../../../common/plugin/postprocess/lib/libYolov8SegmentPostPlugin.so",
        help="yolov8 seg post process plugin path",
    )
    parser.add_argument(
        "--osd_plugin",
        metavar="osd_plugin",
        type=str,
        default="../../../common/plugin/osd/lib/libOsdPlugin.so",
        help="osd plugin path",
    )

    args = parser.parse_args()

    input_file = args.input_image
    output_file = args.output_image
    model_path = args.model
    device_id = args.device_id
    post_plugin_path = args.post_plugin
    osd_plugin_path = args.osd_plugin

    if not os.path.exists(input_file):
        raise ("error: input jpeg file does not exist.")

    if not is_valid_device(device_id):
        raise ("error: device id is invalid.")

    if not os.path.exists(model_path):
        raise ("error: model does not exist.")

    if not os.path.exists(post_plugin_path):
        raise ("error: post process plugin does not exist.")

    if not os.path.exists(osd_plugin_path):
        raise ("error: osd plugin does not exist.")

    # 创建 context
    context, ret = sdk.lyn_create_context(device_id)
    common.error_check(ret, "lyn_create_context")
    ret = sdk.lyn_register_error_handler(common.default_stream_error_handle, None)
    common.error_check(ret, "lyn_register_error_handler")
    # 创建 stream
    stream, ret = sdk.lyn_create_stream()
    common.error_check(ret, "lyn_create_stream")

    # 初始化解码图片类
    print("image decode")

    image_dec = image_decoder.ImageDecoder()
    image_dec.init(input_file, False)
    img_info = image_dec.get_image_dec_info()

    # 同步解码图片到 Device 侧
    decode_img = sdk.lyn_codec_buf_t()
    decode_img.size = img_info.output.predict_buf_size
    p_decode_buf_out, ret = sdk.lyn_malloc(decode_img.size)
    error_check(ret)
    decode_img.data = p_decode_buf_out
    decodeImg_data_ptr = pythonapi.PyCapsule_GetPointer(decode_img.data, None)
    image_dec.decode_image_to_device(stream, decode_img.data)

    # 加载模型
    model = ModelInfo(model_path)

    # 加载 Plugin
    osd_plugin, ret = sdk.lyn_plugin_register(osd_plugin_path)
    common.error_check(ret, "lyn_plugin_register")

    handle = ctypes.CDLL(post_plugin_path)
    if not handle:
        print("error: post plugin load fail.")
        return

    handle.deleteMaskData.argtypes = [ctypes.POINTER(LynSegBoxesInfo)]
    handle.deleteMaskData.restype = ctypes.c_int

    handle.lynSegmentPostProcess.argtypes = [ctypes.POINTER(SegmentPostProcessInfo_t)]
    handle.lynSegmentPostProcess.restype = ctypes.c_int

    segmentPostFunc = handle.lynSegmentPostProcess
    deleteMaskDataFunc = handle.deleteMaskData

    if not segmentPostFunc or not deleteMaskDataFunc:
        print("error: post plugin function err.")
        handle.close()
        return

    # IPE
    print("ipe")
    ipe_buf_apu, ret = sdk.lyn_malloc(model.input_size)
    error_check(ret)
    ipeYolov8 = ipeParamYolov8.IpeParamYolov8(model.width, model.height)

    ipeYolov8.set_img_info(
        img_info.output.width,
        img_info.output.height,
        sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
    )
    ipeYolov8.calc_param(stream, decode_img.data, ipe_buf_apu)

    # 调用 APU 推理接口
    print("infer")
    infer_apu_ptr, ret = sdk.lyn_malloc(model.output_size)
    error_check(ret)
    error_check(
        sdk.lyn_execute_model_async(
            stream,
            model.model,
            ipe_buf_apu,
            infer_apu_ptr,
            model.batch_size,
        )
    )

    infer_host_np = np.zeros(model.output_size, dtype=np.uint8)
    infer_host_ptr = sdk.lyn_numpy_to_ptr(infer_host_np)
    infer_host_ptr_c = pythonapi.PyCapsule_GetPointer(infer_host_ptr, None)

    error_check(
        sdk.lyn_memcpy_async(
            stream,
            infer_host_ptr,
            infer_apu_ptr,
            model.output_size,
            sdk.lyn_memcpy_dir_t.ServerToClient,
        )
    )

    error_check(sdk.lyn_synchronize_stream(stream))

    seg_box_host_np = np.zeros(ctypes.sizeof(LynSegBoxesInfo), dtype=np.uint8)
    seg_box_host_ptr = sdk.lyn_numpy_to_ptr(seg_box_host_np)
    seg_box_host_ptr_c = pythonapi.PyCapsule_GetPointer(seg_box_host_ptr, None)

    print("post")
    post_para = SegmentPostProcessInfo_t()
    post_para.model_width = model.width
    post_para.model_height = model.height
    post_para.img_width = img_info.output.width
    post_para.img_height = img_info.output.height
    post_para.padx = ipeYolov8.m_iPadX
    post_para.pady = ipeYolov8.m_iPadY
    post_para.fatio = ipeYolov8.m_fRatio
    post_para.output_tensor = infer_host_ptr_c
    post_para.boxesInfo = seg_box_host_ptr_c

    segmentPostFunc(post_para)

    seg_box_apu_ptr, ret = sdk.lyn_malloc(ctypes.sizeof(LynSegBoxesInfo))
    error_check(ret)
    error_check(
        sdk.lyn_memcpy(
            seg_box_apu_ptr,
            seg_box_host_ptr,
            ctypes.sizeof(LynSegBoxesInfo),
            sdk.lyn_memcpy_dir_t.ClientToServer,
        )
    )
    seg_box_apu_ptr_c = pythonapi.PyCapsule_GetPointer(seg_box_apu_ptr, None)
    seg_box_info = ctypes.cast(
        seg_box_host_ptr_c, ctypes.POINTER(LynSegBoxesInfo)
    ).contents

    has_mask = False
    if seg_box_info.maskDataLen > 0:
        has_mask = True
        dev_mask_data, ret = sdk.lyn_malloc(seg_box_info.maskDataLen)
        error_check(ret)

        buffer = ctypes.cast(
            seg_box_info.maskData,
            ctypes.POINTER(ctypes.c_uint8 * seg_box_info.maskDataLen),
        ).contents

        mask_data_np = np.frombuffer(
            buffer,
            dtype=np.uint8,
        )

        mask_data_ptr_host = sdk.lyn_numpy_to_ptr(mask_data_np)

        error_check(
            sdk.lyn_memcpy_async(
                stream,
                dev_mask_data,
                mask_data_ptr_host,
                seg_box_info.maskDataLen,
                sdk.lyn_memcpy_dir_t.ClientToServer,
            )
        )

        deleteMaskDataFunc(seg_box_info)

        print("draw")

        osd_para = struct.pack(
            "1P3i1P4i",
            seg_box_apu_ptr_c,
            img_info.output.width,
            img_info.output.height,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
            decodeImg_data_ptr,
            24,
            4,
            2,
            4,
        )
        error_check(
            sdk.lyn_plugin_run_async(
                stream, osd_plugin, "lynDrawSegBoxAndText", osd_para, len(osd_para)
            )
        )

        error_check(sdk.lyn_synchronize_stream(stream))

        imgBgrSize = img_info.output.width * img_info.output.height * 3
        img_bgr_p, ret = sdk.lyn_malloc(imgBgrSize)
        error_check(ret)
        ipe_bgr = ipeParamYolov8.IpeParamC2C(sdk.lyn_pixel_format_t.LYN_PIX_FMT_BGR24)
        ipe_bgr.set_img_info(
            img_info.output.width,
            img_info.output.height,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
        )
        ipe_bgr.calc_param(stream, decode_img.data, img_bgr_p)

        img_bgr_p_c = pythonapi.PyCapsule_GetPointer(img_bgr_p, None)
        dev_mask_data_c = pythonapi.PyCapsule_GetPointer(dev_mask_data, None)

        osd_mask_para = struct.pack(
            "1P3i1f2P1i",
            img_bgr_p_c,
            img_info.output.width,
            img_info.output.height,
            3,
            0.5,
            seg_box_apu_ptr_c,
            dev_mask_data_c,
            seg_box_info.maskDataLen,
        )

        error_check(
            sdk.lyn_plugin_run_async(
                stream, osd_plugin, "maskDraw", osd_mask_para, len(osd_mask_para)
            )
        )

        ipe_nv12 = ipeParamYolov8.IpeParamC2C(sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12)
        ipe_nv12.set_img_info(
            img_info.output.width,
            img_info.output.height,
            sdk.lyn_pixel_format_t.LYN_PIX_FMT_BGR24,
        )
        ipe_nv12.calc_param(stream, img_bgr_p, decode_img.data)

    error_check(sdk.lyn_synchronize_stream(stream))

    print("save new image")
    encoder = image_encoder.ImageEncoder()
    encoder.encode_image_and_save(stream, output_file, decode_img, img_info)

    ipeYolov8.destory()
    if has_mask:
        ipe_bgr.destory()
        ipe_nv12.destory()
        error_check(sdk.lyn_free(dev_mask_data))
        error_check(sdk.lyn_free(img_bgr_p))

    # 等待流处理完
    ret = sdk.lyn_synchronize_stream(stream)
    common.error_check(ret, "lyn_synchronize_stream")
    ret = sdk.lyn_destroy_stream(stream)
    common.error_check(ret, "lyn_destroy_stream")

    error_check(sdk.lyn_free(decode_img.data))
    error_check(sdk.lyn_free(ipe_buf_apu))
    error_check(sdk.lyn_free(infer_apu_ptr))
    error_check(sdk.lyn_free(seg_box_apu_ptr))

    # 卸载模型
    ret = sdk.lyn_plugin_unregister(osd_plugin)
    common.error_check(ret, "lyn_plugin_unregister")
    model.unload_model()
    # 销毁 context
    ret = sdk.lyn_destroy_context(context)
    common.error_check(ret, "lyn_destroy_context")
    handle.close()


if __name__ == "__main__":
    main()
